package ratelimit

import (
	_ "embed"
	"fmt"
	"gitee.com/xiao_hange/go-admin-pkg/pkg/ratelimit"
	"github.com/gin-gonic/gin"
	"net/http"
)

type Builder struct {
	prefix  string
	limiter ratelimit.Limiter
}

func NewBuilder(limiter ratelimit.Limiter) *Builder {
	return &Builder{
		prefix:  "ip-limiter",
		limiter: limiter,
	}
}

func (b *Builder) Prefix(prefix string) *Builder {
	b.prefix = prefix
	return b
}

func (b *Builder) Build() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		limited, err := b.limit(ctx)
		type Result struct {
			Code int    `json:"code"`
			Msg  string `json:"msg"`
		}
		if err != nil {
			ctx.JSON(http.StatusOK, Result{
				Code: http.StatusInternalServerError,
				Msg:  "服务器内部错误",
			})
			return
		}
		if limited {
			ctx.JSON(http.StatusOK, Result{
				Code: http.StatusTooManyRequests,
				Msg:  "已被限流",
			})
			ctx.Abort()
			return
		}
		ctx.Next()
	}
}

func (b *Builder) limit(ctx *gin.Context) (bool, error) {
	key := fmt.Sprintf("%s:%s", b.prefix, ctx.ClientIP())
	return b.limiter.Limited(ctx, key)
}
